import numpy as np
import pandas as pd
from pyecharts import options as opts
from pyecharts.charts import Map, Timeline

def process_data(data_base, data):
    dic_db = dict(data_base)
    dic_d = dict(data)
    
    for k in dic_d:
        dic_db[k] = dic_d[k]

    data_base = list(zip(dict(dic_db).keys(), dict(dic_db).values()))
    return data_base

def get_data_base():
    data_base = [
            ['Alabama', 0],
            ['Alaska', 0],
            ['Arizona', 0],
            ['Arkansas', 0],
            ['California', 0],
            ['Colorado', 0],
            ['Connecticut', 0],
            ['Delaware', 0],
            ['District of Columbia', 0],
            ['Florida', 0],
            ['Georgia', 0],
            ['Hawaii', 0],
            ['Idaho', 0],
            ['Illinois', 0],
            ['Indiana', 0],
            ['Iowa', 0],
            ['Kansas', 0],
            ['Kentucky', 0],
            ['Louisiana', 0],
            ['Maine', 0],
            ['Maryland', 0],
            ['Massachusetts', 0],
            ['Michigan', 0],
            ['Minnesota', 0],
            ['Mississippi', 0],
            ['Missouri', 0],
            ['Montana', 0],
            ['Nebraska', 0],
            ['Nevada', 0],
            ['New Hampshire', 0],
            ['New Jersey', 0],
            ['New Mexico', 0],
            ['New York', 0],
            ['North Carolina', 0],
            ['North Dakota', 0],
            ['Ohio', 0],
            ['Oklahoma', 0],
            ['Oregon', 0],
            ['Pennsylvania', 0],
            ['Rhode Island', 0],
            ['South Carolina', 0],
            ['South Dakota', 0],
            ['Tennessee', 0],
            ['Texas', 0],
            ['Utah', 0],
            ['Vermont', 0],
            ['Virginia', 0],
            ['Washington', 0],
            ['West Virginia', 0],
            ['Wisconsin', 0],
            ['Wyoming', 0],
            ['Puerto Rico', 0]
        ]

    return data_base

def draw():
    data_set = pd.read_csv('./web/usa/us-states.csv')
    data_base = get_data_base()

    high = data_set.max()['cases']
    low = 0

    pieces_range = np.logspace(0, np.log10(high)+1, 8)
    pieces = [{'min': pieces_range[i], 'max': pieces_range[i+1], 'label':'1e{}'.format(i+1)} for i in range(len(pieces_range)-1)]
    pieces.append({'max': 0, 'max': 1, 'label':'0'})

    num = 0

    tl = Timeline()
    for i in data_set.groupby('date'):
        if(num == 20):
            break
        data = i[1][['state', 'cases']].values
        data_base = process_data(data_base, data)
        c = (
            Map()
            .add("人数", data_base, "美国")
            .set_global_opts(title_opts=opts.TitleOpts(title="美国新冠疫情热力图"), visualmap_opts=opts.VisualMapOpts(max_=10979626, pieces=pieces, is_piecewise=True, min_=low, split_number=10))
        )
        tl.add(c, "{}年".format(i[0]))
        num += 1

    return tl.render_embed()